import numpy as np
from isaacgym import gymutil
from isaacgym import gymapi
#from isaacgym import gymtorch
from math import sqrt
import torch
import time

headless = True

gym = gymapi.acquire_gym()

# get default set of parameters
sim_params = gymapi.SimParams()

args = gymutil.parse_arguments()

sim_params.physx.solver_type = 1
sim_params.physx.num_position_iterations = 4
sim_params.physx.num_velocity_iterations = 1
sim_params.physx.num_threads = args.num_threads
sim_params.physx.use_gpu = args.use_gpu
sim_params.use_gpu_pipeline = args.use_gpu_pipeline

# set common parameters
sim_params.substeps = 2
sim_params.dt = 1.0 / 60.0

# setting up the Z-up axis
sim_params.up_axis = gymapi.UP_AXIS_Z
sim_params.gravity = gymapi.Vec3(0.0, 0.0, -9.8)

# sim = gym.create_sim(args.compute_device_id, args.graphics_device_id, gymapi.SIM_PHYSX, sim_params)
sim = gym.create_sim(args.compute_device_id, -1, gymapi.SIM_PHYSX, sim_params)

device = args.sim_device if args.use_gpu_pipeline else 'cpu'

# configure the ground plane
plane_params = gymapi.PlaneParams()
plane_params.normal = gymapi.Vec3(0, 0, 1) # z-up!
plane_params.distance = 0
plane_params.static_friction = 1
plane_params.dynamic_friction = 1
plane_params.restitution = 0.5

# create the ground plane
gym.add_ground(sim, plane_params)

# set up the env grid
num_envs = 1
spacing = 0.75
env_lower = gymapi.Vec3(-spacing, 0.0, -spacing)
env_upper = gymapi.Vec3(spacing, spacing, spacing)

# Attractor setup
attractor_properties = gymapi.AttractorProperties()
attractor_properties.stiffness = 5e10
attractor_properties.damping = 0
attractor_properties.axes = gymapi.AXIS_ROTATION
pose_ = gymapi.Transform()
pose_.p = gymapi.Vec3(0, 0.0, 0.0)
pose_.r = gymapi.Quat(0,1,0,1)


# create ant asset
asset_root = "./assets"
asset_file = "mjcf/nv_ant.xml"
print("Loading asset '%s' from '%s'" % (asset_file, asset_root))
ant_asset = gym.load_asset(sim, asset_root, asset_file)
if ant_asset is None:
    raise IOError("Failed to load asset")

# create procedural asset
asset_options = gymapi.AssetOptions()
asset_options.density = 1.0
asset_options.fix_base_link = False

asset_options1 = gymapi.AssetOptions()
asset_options1.density = 1.0
asset_options1.fix_base_link = True

box_asset = gym.create_box(sim, 1, 1, 1, asset_options)
sphere_asset = gym.create_sphere(sim, 1, asset_options1)
capsule_asset = gym.create_capsule(sim, 1, 1, asset_options)

# if not headless:
#     cam_props = gymapi.CameraProperties()
#     viewer = gym.create_viewer(sim, cam_props)

# create environment
env = gym.create_env(sim, env_lower, env_upper, 2)

# # create capsule_asset actor in the environment
# pose = gymapi.Transform()
# pose.p = gymapi.Vec3(0, 0, 10)
# pose.r = gymapi.Quat(0,1,0,1)
# # ahandle = gym.create_actor(env, ant_asset, pose, 'ant', 0, 1)
# # ahandle = gym.create_actor(env, capsule_asset, pose, 'capsule_asset', 0, 1)
# ahandle = gym.create_actor(env, capsule_asset, pose, 'box_asset', 0, 0)

# # create capsule_asset actor in the environment
# pose = gymapi.Transform()
# pose.p = gymapi.Vec3(1.5, 4, 0)
# pose.r = gymapi.Quat(0,1,0,1)
# bhandle = gym.create_actor(env, sphere_asset, pose, 'sphere_asset', 0, 0)

# # set shape physics properties for target boxes
# shape_props = gym.get_actor_rigid_shape_properties(env, ahandle)
# shape_props[0].friction = 0.
# shape_props[0].rolling_friction = 0.
# shape_props[0].torsion_friction = 0.
# shape_props[0].compliance = 1
# shape_props[0].restitution = 1
# gym.set_actor_rigid_shape_properties(env, ahandle, shape_props)
# gym.set_actor_rigid_shape_properties(env, bhandle, shape_props)

# # Initialize the attractor
# attractor_properties.target = pose_
# attractor_properties.rigid_handle = ahandle
# attractor_handle = gym.create_rigid_body_attractor(env, attractor_properties)


# gym.prepare_sim(sim)


# torque_amt = 0
# frame_count = 0
# num_agents = 2

# # time test
# start_time = time.time()
# last_frame_cnt = 0

# while True:

#     if not headless:
#         if gym.query_viewer_has_closed(viewer):
#             break

#     if (frame_count - 99) % 200 == 0:
#         # set forces and torques for the ant root bodies
#         forces = torch.zeros((num_envs, num_agents, 3), device=device, dtype=torch.float)
#         torques = torch.zeros((num_envs, num_agents, 3), device=device, dtype=torch.float)
#         forces[:, :, 1] = 3000
#         torques[:, :, 2] = torque_amt
#         gym.apply_rigid_body_force_tensors(sim, gymtorch.unwrap_tensor(forces), gymtorch.unwrap_tensor(torques), gymapi.ENV_SPACE)

#     # step the physics
#     gym.simulate(sim)
#     gym.fetch_results(sim, True)

#     if not headless:
#         # update the viewer
#         gym.step_graphics(sim)
#         gym.draw_viewer(viewer, sim, True)
#         # Wait for dt to elapse in real time.
#         # This synchronizes the physics simulation with the rendering rate.
#         # gym.sync_frame_time(sim)

#     frame_count += 1
#     end_time = time.time()
#     if (end_time - start_time) > 1:
#         print("FPS: %.2f" % (frame_count - last_frame_cnt))
#         last_frame_cnt = frame_count
#         start_time = time.time()

# if not headless:
#     gym.destroy_viewer(viewer)
#     gym.destroy_sim(sim)